﻿#pragma once

#include <async_simple/coro/Lazy.h>
#include <async_simple/coro/ConditionVariable.h>
#include <asio.hpp>
#include <memory>
#include <semaphore>
#include <lwip/tcp.h>

#include "io.hpp"
#include "coro_io.hpp"
#include "io_context_pool.hpp"
#include "ring_buf.hpp"
#include "netpacket_pool.h"
#include "base/xlog.h"

// #define USE_CORO

namespace driver2socks {

	inline bool parse_udp_proxy_header(const void* buf, std::size_t len,
		asio::ip::tcp::endpoint& src, asio::ip::tcp::endpoint& dst, uint16_t& payload_len)
	{
		const uint8_t* p = (const uint8_t*)buf;
		if (len < 16)
			return false;

		read_int8(p);
		auto local_ip = read_uint32(p);
		auto local_port = read_uint16(p);
		src.address(asio::ip::address_v4(local_ip));
		src.port(local_port);

		read_int8(p);
		auto remote_ip = read_uint32(p);
		auto remote_port = read_uint16(p);
		dst.address(asio::ip::address_v4(remote_ip));
		dst.port(remote_port);

		payload_len = read_uint16(p);

		return true;
	}

	inline std::string make_udp_proxy_header(
		const asio::ip::tcp::endpoint& src, const asio::ip::tcp::endpoint& dst, const uint16_t& payload_len)
	{
		std::string response;
		response.resize(payload_len + 16);
		char* resp = (char*)response.data();

		// 添加头信息.
		write_uint8(1, resp);	// atyp.
		write_uint32(src.address().to_v4().to_uint(), resp);	// ip.
		write_uint16(src.port(), resp);	// port.
		write_uint8(1, resp);	// atyp.
		write_uint32(dst.address().to_v4().to_uint(), resp);	// ip.
		write_uint16(dst.port(), resp);	// port.
		write_uint16(payload_len, resp); // payload_len.

		return response;
	}


	//////////////////////////////////////////////////////////////////////////

	class error_category_impl;

	template<class error_category>
    const std::error_category& error_category_single()
	{
		static error_category error_category_instance;
        return reinterpret_cast<const std::error_category&>(error_category_instance);
	}

    inline const std::error_category& error_category()
	{
		return error_category_single<driver2socks::error_category_impl>();
	}

	namespace errc {
		enum errc_t
		{
			/// SOCKS unsupported version.
			socks_unsupported_version = 1000,

			/// SOCKS username required.
			socks_username_required,

			/// SOCKS unsupported authentication version.
			socks_unsupported_authentication_version,

			/// SOCKS authentication error.
			socks_authentication_error,

			/// SOCKS general failure.
			socks_general_failure,

			/// SOCKS command not supported.
			socks_command_not_supported,

			/// SOCKS no identd running.
			socks_no_identd,

			/// SOCKS no identd running.
			socks_identd_error,

			/// request rejected or failed.
			socks_request_rejected_or_failed,

			/// request rejected becasue SOCKS server cannot connect to identd on the client.
			socks_request_rejected_cannot_connect,

			/// request rejected because the client program and identd report different user - ids
			socks_request_rejected_incorrect_userid,

			socks_connect_proxy_fail,
		};

        inline std::error_code make_error_code(errc_t e)
		{
            return std::error_code(static_cast<int>(e), driver2socks::error_category());
		}
	}

    class error_category_impl : public std::error_category
	{
		virtual const char* name() const noexcept
		{
			return "SOCKS";
		}

		virtual std::string message(int e) const
		{
			switch (e)
			{
			case errc::socks_unsupported_version:
				return "SOCKS unsupported version";
			case errc::socks_username_required:
				return "SOCKS username required";
			case errc::socks_unsupported_authentication_version:
				return "SOCKS unsupported authentication version";
			case errc::socks_authentication_error:
				return "SOCKS authentication error";
			case errc::socks_general_failure:
				return "SOCKS general failure";
			case errc::socks_command_not_supported:
				return "SOCKS command not supported";
			case errc::socks_no_identd:
				return "SOCKS no identd running";
			case errc::socks_identd_error:
				return "SOCKS identd error";
			case errc::socks_request_rejected_or_failed:
				return "request rejected or failed";
			case errc::socks_request_rejected_cannot_connect:
				return "request rejected becasue SOCKS server cannot connect to identd on the client";
			case errc::socks_request_rejected_incorrect_userid:
				return "request rejected because the client program and identd report different user";
			case errc::socks_connect_proxy_fail:
				return "proxy connect fail";
			default:
				return "Unknown PROXY error";
			}
		}
	};
}

namespace std {

	template <>
	struct is_error_code_enum<driver2socks::errc::errc_t>
	{
		static const bool value = true;
	};

} // namespace system


namespace driver2socks {

	//////////////////////////////////////////////////////////////////////////

	// 解析uri格式
	// scheme:[//[user[:password]@]host[:port]][/path][?query][#fragment]

    struct SocksAddress
	{
        typedef std::shared_ptr<SocksAddress> Ptr;
		std::string host;
        std::string port;
		std::string query;
		std::string fragment;
		std::string username;
		std::string password;

		// proxy_address为代理服务器连接的对象,
		// 如果是ip, proxy_hostname则应该为false.
		// 如果是域名, proxy_hostname应该为true.
		std::string proxy_address;

		// 代理服务器连接的目标端口.
		std::string proxy_port;

		// 控制代理服务器是否解析域名.
		bool proxy_hostname;

		// 打开udp转发.
		bool udp_associate;
	};

    class SocksClient
	{
    public:

		enum {
			SOCKS_VERSION_4 = 4,
			SOCKS_VERSION_5 = 5
		};
		enum {
			SOCKS5_AUTH_NONE = 0x00,
			SOCKS5_AUTH = 0x02,
			SOCKS5_AUTH_UNACCEPTABLE = 0xFF
		};
		enum {
			SOCKS_CMD_CONNECT = 0x01,
			SOCKS_CMD_BIND = 0x02,
			SOCKS5_CMD_UDP = 0x03
		};
		enum {
			SOCKS5_ATYP_IPV4 = 0x01,
			SOCKS5_ATYP_DOMAINNAME = 0x03,
			SOCKS5_ATYP_IPV6 = 0x04
		};
		enum {
			SOCKS5_SUCCEEDED = 0x00,
			SOCKS5_GENERAL_SOCKS_SERVER_FAILURE,
			SOCKS5_CONNECTION_NOT_ALLOWED_BY_RULESET,
			SOCKS5_NETWORK_UNREACHABLE,
			SOCKS5_CONNECTION_REFUSED,
			SOCKS5_TTL_EXPIRED,
			SOCKS5_COMMAND_NOT_SUPPORTED,
			SOCKS5_ADDRESS_TYPE_NOT_SUPPORTED,
			SOCKS5_UNASSIGNED
		};
		enum {
			SOCKS4_REQUEST_GRANTED = 90,
			SOCKS4_REQUEST_REJECTED_OR_FAILED,
			SOCKS4_CANNOT_CONNECT_TARGET_SERVER,
			SOCKS4_REQUEST_REJECTED_USER_NO_ALLOW,
		};

		enum {
			MAX_RECV_BUFFER_SIZE = 768,	// 最大udp接收缓冲大小.
			MAX_SEND_BUFFER_SIZE = 768	// 最大udp发送缓冲大小.
		};

	public:
        explicit SocksClient():
            executor_(coro_io::get_global_executor()),
            socket_(executor_->get_asio_executor()),
            buf_send_(128)

		{
		}
        ~SocksClient()
		{
            closeSocket();
		}
        void setRecvFunc(std::function<void (NetPacket::Ptr)> f)
        {
            recv_func_ = f;
        }
        void setErrorFunc(std::function<void (std::error_code&)> f)
        {
            error_func_ = f;
        }
        bool sendData(NetPacket::Ptr buff)
        {
            if (is_stop_.test()) return false;
#ifndef USE_CORO
            bool ret = buf_send_.try_emplace(buff);
            asio::post(executor_->get_asio_executor(),std::bind(&SocksClient::startSend,this));
            return ret;
#else
            return buf_send_.Write(data, len);
#endif

		}
		void closeSocket()
		{
            if (is_stop_.test_and_set()) return;
            asio::steady_timer timer(executor_->get_asio_executor());
            timer.expires_after(std::chrono::milliseconds(300));
            timer.async_wait([this](asio::error_code ec){
                if (ec == asio::error::operation_aborted) return;
                is_writing_.clear();
                is_writing_.notify_all();
                is_reading_.clear();
                is_reading_.notify_all();
                std::cerr << "close socksclient timeout!\n";
            });

            SDEBUG("Proxy closeSocket,{},port:{}",remote_address_,remote_port_);
            asio::error_code ignore_ec;
            socket_.close(ignore_ec);
            // auto begin = std::chrono::high_resolution_clock::now();
            is_writing_.wait(true);
            is_reading_.wait(true);
            // auto end = std::chrono::high_resolution_clock::now();
            // std::cerr << "wait time:"<<std::chrono::duration_cast<std::chrono::milliseconds>(end-begin)<<"\n";
            timer.cancel();
		}
#ifndef USE_CORO
		void startSend()
        {
            if (buf_send_.empty()) {
                is_writing_.clear();
                if (is_stop_.test()) {
                    is_writing_.notify_all();
                }
				return;
			}

            if (is_writing_.test_and_set(std::memory_order_acquire)) return;

            NetPacket::Ptr buff = *(buf_send_.front());
            buf_send_.pop();
            asio::async_write(socket_, asio::buffer(buff->data,buff->data_len), asio::transfer_exactly(buff->data_len),
                [this, buff](asio::error_code ec, size_t size) {
                    is_writing_.clear();
                    if (is_stop_.test()) {
                        is_writing_.notify_all();
                    } else {
                        if (buff->data_len != size) {
                            SERROR("socks_client socket write error,data_len != size");
                        }
                        if (ec) {
                            SERROR("socks_client socket write error:{}",ec.message());
                            error_func_(ec);
                        } else {
                            startSend();
                        }
                    }
			});
		}
        void startRecv()
        {
            is_reading_.test_and_set();
            std::shared_ptr<NetPacket> bf = _NetPacketPool->getSharedPacket(NETPACKET_DATA_SIZE);
            socket_.async_read_some(asio::buffer(bf->data, bf->data_len),
                [bf,this](asio::error_code ec, size_t size) {
                    is_reading_.clear();
                    if (is_stop_.test()) {
                        is_reading_.notify_all();
                        return;
                    }
                    if (ec) {
                        if (ec == asio::error::eof) {
                            SDEBUG("socket eof,{}",socket_.remote_endpoint().address().to_string());
                        } else {
                            SERROR("socks_client socket read error:{}.",ec.message());
                        }
                        error_func_(ec);
                    } else {
                        bf->data_len = size;
                        recv_func_(bf);
                        startRecv();
                    }

                });
        }
#else
        async_simple::coro::Lazy<bool> startSend()
        {
            // while (true) {
            int len = buf_send_.GetAvailable();
            std::cerr << "GetAvailable:"<<len<<"\n";
            if (len > 1024) {
                len = 1024;
            }
            if (len <= 0) {
                co_await async_simple::coro::Yield{};
                co_return true;
            }

            is_writing_.test_and_set();
            std::shared_ptr<NetPacket> bf = _NetPacketPool->getSharedPacket(len);
            buf_send_.Read(bf->data, len);
            auto [ec,size] = co_await coro_io::async_write(socket_, asio::buffer(bf->data,bf->data_len));
            std::cerr << "async_write:"<<size<<"\n";
            is_writing_.clear();
            if (is_close_) {
                is_writing_.notify_all();
                co_return false;
            }
            if (ec) {
                if (!is_close_) error_func_(ec);
                co_return false;
            }

            if (size != len) {
                SERROR("socksclient async_write size != len.");
                co_return false;
            }
            // }
            co_return true;
        }
        async_simple::coro::Lazy<bool> startRecv()
        {
            std::shared_ptr<NetPacket> bf = _NetPacketPool->getSharedPacket(NETPACKET_DATA_SIZE);
            is_reading.test_and_set();
            auto [ec,size] = co_await coro_io::async_read_some(socket_, asio::buffer(bf->data, bf->capacity_size));
            std::cerr << "async_read_some:"<<size<<"\n";
            is_reading.clear();
            if (is_close_) {
                is_reading.notify_all();
                co_return false;
            }
            if (ec) {
                if (!is_close_) error_func_(ec);
                co_return false;
            } else {
                bf->data_len = size;
                recv_func_(bf);
            }

            co_return true;
        }
#endif
        void start_socks(std::string host, uint16_t port,
                         std::string proxy_address, uint16_t proxy_port)
		{
            SocksAddress::Ptr socks_addr = std::make_shared<SocksAddress>();
			socks_addr->host = host;
			socks_addr->port = std::to_string(port);
			socks_addr->proxy_hostname = false;
			socks_addr->udp_associate = false;
			socks_addr->proxy_address = proxy_address;
			socks_addr->proxy_port = std::to_string(proxy_port);
            socks_address_ = socks_addr;

            remote_address_ = proxy_address;
            remote_port_ = proxy_port;
/*
            is_writing_.test_and_set();
            async_do_proxy().setEx(executor_).start([this](async_simple::Try<void> Result){
                if (Result.hasError()) {
                    SERROR("Error Happened in async_do_proxy.");
                } else {
#ifndef USE_CORO
                    is_writing_.clear();
                    startSend();
                    startRecv();
#endif
                }
            });
*/
            is_writing_.test_and_set(std::memory_order_acquire);
            async_do_proxy().start([this](async_simple::Try<void> Result){
                if (Result.hasError()) {
                    SERROR("Error Happened in async_do_proxy.");
                } else {
#ifndef USE_CORO
                    is_writing_.clear();
                    startSend();
                    startRecv();
#endif
                }
            });
		}

        async_simple::coro::Lazy<bool> async_do_proxy()
		{

            auto ec = co_await coro_io::async_connect(executor_,socket_,socks_address_->host,socks_address_->port);
            if (ec) {
                SERROR("Proxy connect failed:{}",ec.message());
                std::error_code error_code = errc::socks_connect_proxy_fail;
                error_func_(error_code);
                co_return false;
			}

            ec = co_await do_socks5();
            if (ec) {
                SERROR("Proxy handshake failed.");
            } else {
                SDEBUG("Proxy server handshake complete,{},port:{}",remote_address_,remote_port_);
            }

#ifdef USE_CORO
            for (bool ret = true;ret;) {
                co_await startSend();
                co_await startRecv();
                // co_await async_simple::coro::Yield{};
            }
#endif
			co_return true;
        }

		asio::ip::udp::endpoint udp_endpoint()
		{
			asio::ip::udp::endpoint endp;
			endp.address(m_remote_endp.address());
			endp.port(m_remote_endp.port());
			return endp;
		}

	private:
        void onError(std::error_code& ec)
        {
            error_func_(ec);
        }

        async_simple::coro::Lazy<std::error_code> do_socks5()
		{
            std::size_t bytes_to_write = socks_address_->username.empty() ? 3 : 4;
			asio::streambuf request;
			asio::mutable_buffer b = request.prepare(bytes_to_write);
			char* p = asio::buffer_cast<char*>(b);

			write_uint8(5, p); // SOCKS VERSION 5.
            if (socks_address_->username.empty())
			{
				write_uint8(1, p); // 1 authentication method (no auth)
				write_uint8(0, p); // no authentication
			}
			else
			{
				write_uint8(2, p); // 2 authentication methods
				write_uint8(0, p); // no authentication
				write_uint8(2, p); // username/password
			}
			
			request.commit(bytes_to_write);
			//发起握手请求
            auto [ec1, write_size] = co_await coro_io::async_write(socket_, request.data());
			
            if (ec1) {
                co_return ec1;
			}

			asio::streambuf response;
            auto [ec, read_size] = co_await coro_io::async_read(socket_, response, 2);
            if (ec) {
                co_return ec;
			}
			
			//处理代理服务器返回的握手响应
			int method;
			bool authed = false;

			{
				int version;

				asio::const_buffer b = response.data();
				const char* p = asio::buffer_cast<const char*>(b);
				version = read_uint8(p);
				method = read_uint8(p);
				if (version != 5)	// 版本不等于5, 不支持socks5.
				{
                    ec = driver2socks::errc::socks_unsupported_version;
                    co_return ec;
				}
			}
			//如果代理服务需要密码验证
			if (method == 2)
			{
                if (socks_address_->username.empty())
				{
                    ec = driver2socks::errc::socks_username_required;
                    co_return ec;
				}

				// start sub-negotiation.
				request.consume(request.size());

                std::size_t bytes_to_write = socks_address_->username.size() + socks_address_->password.size() + 3;
				asio::mutable_buffer mb = request.prepare(bytes_to_write);
				char* mp = asio::buffer_cast<char*>(mb);

				write_uint8(1, mp);
                write_uint8(static_cast<int8_t>(socks_address_->username.size()), mp);
                write_string(socks_address_->username, mp);
                write_uint8(static_cast<int8_t>(socks_address_->password.size()), mp);
                write_string(socks_address_->password, mp);
				request.commit(bytes_to_write);

				// 发送用户密码信息.
                auto [ec, wsize] = co_await coro_io::async_write(socket_, request.data());
                if (ec) {
                    co_return ec;
				}
				//BOOST_ASSERT("len == bytes_to_write" && len == bytes_to_write);

				// 读取状态.
				response.consume(response.size());
                auto [ec2, rsize] = co_await coro_io::async_read(socket_, response, 2);
                if (ec2) {
                    co_return ec2;
				}
				//BOOST_ASSERT("len == 2" && len == 2);

				// 读取版本状态.
				asio::const_buffer cb = response.data();
				const char* cp = asio::buffer_cast<const char*>(cb);

				int version = read_uint8(cp);
				int status = read_uint8(cp);

				// 不支持的认证版本.
				if (version != 1)
				{
                    ec = errc::socks_unsupported_authentication_version;
                    co_return ec;
				}

				// 认证错误.
				if (status != 0)
				{
                    ec = errc::socks_authentication_error;
                    co_return ec;
				}

				authed = true;
			}//密码认证结束

			//不需要认证
			if (method == 0 || authed)
			{
				request.consume(request.size());
                std::size_t bytes_to_write = 7 + remote_address_.size();
				asio::mutable_buffer mb = request.prepare(std::max<std::size_t>(bytes_to_write, 22));
				char* wp = asio::buffer_cast<char*>(mb);
				
				// 发送socks5连接命令.
				write_uint8(5, wp); // SOCKS VERSION 5.
									// CONNECT/UDP command.
                write_uint8(socks_address_->udp_associate ? SOCKS5_CMD_UDP : SOCKS_CMD_CONNECT, wp);
				write_uint8(0, wp); // reserved.

                if (socks_address_->proxy_hostname)
				{
                    write_uint8(3, wp); // atyp, domain name.
                    write_uint8(static_cast<int8_t>(remote_address_.size()), wp);	// domainname size.
                    std::copy(remote_address_.begin(), remote_address_.end(), wp);		// domainname.
                    wp += remote_address_.size();
                    write_uint16(remote_port_, wp);					// port.
				}
				else
				{
                    auto endp = asio::ip::address::from_string(remote_address_);
					if (endp.is_v4())
					{
						write_uint8(1, wp); // ipv4.
						write_uint32(endp.to_v4().to_ulong(), wp);
                        write_uint16(remote_port_, wp);
						bytes_to_write = 10;
					}
					else
					{
						write_uint8(4, wp); // ipv6.
						auto bytes = endp.to_v6().to_bytes();
						std::copy(bytes.begin(), bytes.end(), wp);
						wp += 16;
                        write_uint16(remote_port_, wp);
						bytes_to_write = 22;
					}
				}

				std::size_t len = 0;
				request.commit(bytes_to_write);
                auto [ec, wsize] = co_await coro_io::async_write(socket_, request.data());
                if (ec) {
                    co_return ec;
				}
				//BOOST_ASSERT("len == bytes_to_write" && len == bytes_to_write);
				
				//读取连接响应
				std::size_t bytes_to_read = 10;
				response.consume(response.size());
                auto [ec3, rsize] = co_await coro_io::async_read(socket_, response, bytes_to_read);
                if (ec3) {
                    co_return ec3;
				}
				//BOOST_ASSERT("len == bytes_to_read" && len == bytes_to_read);
				asio::const_buffer cb = response.data();
				const char* rp = asio::buffer_cast<const char*>(cb);
				int version = read_uint8(rp);
				int resp = read_uint8(rp);
				read_uint8(rp);	// skip RSV.
				int atyp = read_uint8(rp);

				if (atyp == 1) // ADDR.PORT
				{
					m_remote_endp.address(asio::ip::address_v4(read_uint32(rp)));
					m_remote_endp.port(read_uint16(rp));

                    if (socks_address_->udp_associate)
					{
						// 更新远程地址, 后面用于udp传输.
                        m_remote_endp.address(socket_.remote_endpoint(ec).address());
						//LOG_DBG << "* SOCKS udp server: " << m_remote_endp.address().to_string()
						//	<< ":" << m_remote_endp.port();
						// 在这之后，保持这个tcp连接直到udp代理也不需要了.
					}
// 					else
// 					{
// 						LOG_DBG << "* SOCKS remote host: " << m_remote_endp.address().to_string()
// 							<< ":" << m_remote_endp.port();
// 					}

					//response.consume(len);
					//BOOST_ASSERT("response.size() == 0" && response.size() == 0);
				}
				else if (atyp == 3) // DOMAIN
				{
					auto domain_length = read_uint8(rp);

                    auto [ec, rsize] = co_await coro_io::async_read(socket_, response, domain_length - 3);
                    if (ec) {
                        co_return ec;
					}
					//BOOST_ASSERT("len == domain_length - 3" && len == domain_length - 3);

					asio::const_buffer cb = response.data();
					rp = asio::buffer_cast<const char*>(cb) + 5;

					std::string domain;
					for (int i = 0; i < domain_length; i++)
						domain.push_back(read_uint8(rp));
					auto port = read_uint16(rp);

					//LOG_DBG << "* SOCKS remote host: " << domain << ":" << port;
					response.consume(len + 10);
					//BOOST_ASSERT("response.size() == 0" && response.size() == 0);
                } else {
                    ec = errc::socks_general_failure;
                    co_return ec;
				}

                if (version != 5) {
                    ec = errc::socks_unsupported_version;
                    co_return ec;
				}

                if (resp != 0) {
					ec = errc::socks_general_failure;
					// 得到更详细的错误信息.
					switch (resp)
					{
					case 2: ec = asio::error::no_permission; break;
					case 3: ec = asio::error::network_unreachable; break;
					case 4: ec = asio::error::host_unreachable; break;
					case 5: ec = asio::error::connection_refused; break;
					case 6: ec = asio::error::timed_out; break;
					case 7: ec = errc::socks_command_not_supported; break;
					case 8: ec = asio::error::address_family_not_supported; break;
					}

                    co_return ec;
				}
				
                ec = asio::error_code();	// 没有发生错误, 返回.
                co_return ec;
			}

            ec = asio::error::address_family_not_supported;
            co_return ec;
		}

	private:
        coro_io::ExecutorWrapper<> *executor_;
        std::function<void (NetPacket::Ptr)> recv_func_;
        std::function<void (std::error_code&)> error_func_;

        std::atomic_flag is_writing_ = ATOMIC_FLAG_INIT;
        std::atomic_flag is_reading_ = ATOMIC_FLAG_INIT;
        std::atomic_flag is_stop_ = ATOMIC_FLAG_INIT;

        asio::ip::tcp::socket socket_;
        SocksAddress::Ptr socks_address_;
        std::string remote_address_;
        uint16_t remote_port_;
		asio::ip::tcp::endpoint m_remote_endp;

        SPSCQueue<NetPacket::Ptr> buf_send_;
	};
}

